
#include "Home.h"
#include "Util.h"
#include "ProDataIO.h"
#include "Fit.h"

#ifdef GSL
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>

int DoFit(const gsl_multifit_robust_type *T, 
          const gsl_matrix *X, 
          const gsl_vector *y, gsl_vector *c, 
          gsl_matrix *cov)
{
  int s;
  gsl_multifit_robust_workspace * work
    = gsl_multifit_robust_alloc (T, X->size1, X->size2);

  s = gsl_multifit_robust (X, y, c, cov, work);
  gsl_multifit_robust_free (work);

  return s;
}

void RobustMultiFit(const Table_t &table, 
                    const int xId, const int yId, // index in the table
                    double &c0, double &c1 /* out parameters */, Table_t &fitTable)
{
#ifdef DEBUG
    if(table.variables.size() >= xID || table.varibales.siee() >= yID)
        Fatal("%s: xId=%d yId=%d",__func__, xId, yId);
#endif

    int numPoints = table.data.size();

    gsl_matrix *X   = gsl_matrix_alloc(numPoints,2);
    gsl_vector *y   = gsl_vector_alloc(numPoints);
    gsl_vector *c   = gsl_vector_alloc(2);
    gsl_matrix *cov = gsl_matrix_alloc(2,2);

    //consrtuct
    double minX=0, maxX=0, xValue;
    for(auto i=0; i<table.data.size(); i++){
        xValue = table.data[i][xId];

        if(xValue<minX){
            minX=xValue;
        }else if(xValue>maxX){
            maxX = xValue;
        }

        gsl_matrix_set(X, i, 0, 1.0);
        gsl_matrix_set(X, i, 1, xValue);
        gsl_vector_set(y, i, table.data[i][yId]);
//      printf("%g %g %g\n",gsl_matrix_get(X,i,0),gsl_matrix_get(X,i,1),gsl_vector_get(y,i));
    }

    DoFit(gsl_multifit_robust_bisquare, X, y, c, cov);

    c0 = gsl_vector_get(c,0);
    c1 = gsl_vector_get(c,1);

    InitTable(fitTable);

    fitTable.aux["c0"] = std::to_string(c0);
    fitTable.aux["c1"] = std::to_string(c1);

    fitTable.aux["cov00"] = std::to_string(gsl_matrix_get(cov,0,0));
    fitTable.aux["cov01"] = std::to_string(gsl_matrix_get(cov,0,1));
    fitTable.aux["cov10"] = std::to_string(gsl_matrix_get(cov,1,0));
    fitTable.aux["cov11"] = std::to_string(gsl_matrix_get(cov,1,1));

    const int nPoints = 100;
    double length=maxX-minX, centor=0.5*(maxX+minX);
    maxX = centor + 0.6*length;
    minX = centor - 0.6*length;
    double stride = 1.2*length/(double(nPoints)-1.0);
    
    fitTable.variables.push_back(table.variables[xId]);
    fitTable.variables.push_back(table.variables[yId]);

    fitTable.data.resize(nPoints);
    for(int i=0; i<nPoints; i++){
        fitTable.data[i].resize(2);
        fitTable.data[i][0] = minX + stride*double(i);
        fitTable.data[i][1] = c0 + c1*fitTable.data[i][0];
    }
    fitTable.i = nPoints;

    gsl_matrix_free(X);
    gsl_matrix_free(cov);
    gsl_vector_free(y);
    gsl_vector_free(c);

}

#endif